#include "field_energy.h"

namespace Field_packing
{
	Field_energy_system::Field_energy_system(const std::vector<Packing_polygon*>& polygon_set, const Container_box& bndbox) : polygons(polygon_set), System_potential_energy(bndbox) 
	{
		rm = 0.0;
		//a = 0.0;
		for (unsigned int i = 0; i < polygons.size(); i++)
			rm += std::fabs(polygons[i]->area());
		rm /= polygons.size();
		rm = std::sqrt(rm);
		//rm /= 10;
		//	a += std::fabs(polygons[i]->area());
		//a /= polygons.size();
		//lambda = 1.0;
	}
	double Field_energy_system::system_energy_at(const double *const x) 
	{
		double me = molecular_energy_at(x);
		return me;
	}
	double Field_energy_system::point_polygon_squared_distance(const Point_2& p, const Packing_polygon& polygon)
	{
		double dist = std::numeric_limits<double>::max();
		for (Packing_polygon::Edge_const_iterator eit = polygon.edges_begin(); eit != polygon.edges_end(); ++eit)
		{
			double edge_dist = CGAL::squared_distance(p, *eit);
			dist = std::min(dist, edge_dist);
		}
		return dist;
	}
	double Field_energy_system::point_polygon_squared_distance(const Point_2& p, const Packing_polygon& polygon, unsigned int& edge_idx)
	{
		double dist = std::numeric_limits<double>::max();
		for (unsigned int i = 0; i < polygon.size(); i++)
		{
			double edge_dist = CGAL::squared_distance(p, polygon.edge(i));
			if (dist > edge_dist)
			{
				edge_idx = i;
				dist = edge_dist;
			}
		}
		return dist;
	}
	double Field_energy_system::molecular_energy_at(const double *const x)
	{
		double LJ_energy = 0.0;
		std::vector<Packing_polygon> transformed_polygons;
		transformed_polygons.reserve(polygons.size());

		for (unsigned int i = 0; i < polygons.size(); i++)
		{
			transformed_polygons.push_back(Packing_polygon(polygons[i]->vertices_begin(), polygons[i]->vertices_end()));
			transformed_polygons.back().transform(Parameter(1.0, x[i*3], x[i*3+1], x[i*3+2]));
		}

		for (unsigned int i = 0; i < transformed_polygons.size(); i++)
		{
			for (unsigned int j = 0; j < transformed_polygons.size(); j++)
			{
				if (i == j)		continue;
				for (Packing_polygon::Vertex_const_iterator vit = transformed_polygons[i].vertices_begin(); vit != transformed_polygons[i].vertices_end(); ++vit)
				{
					double d2 = point_polygon_squared_distance(*vit, transformed_polygons[j]);
					double rm2_over_d2;
					if (d2 == 0.0)
						rm2_over_d2 = 1.0e20;
					else
						rm2_over_d2 = rm*rm/d2;
					double rm6_over_d6 = std::pow(rm2_over_d2, 3.0);
					LJ_energy += rm6_over_d6*(rm6_over_d6 - 2);
				}
			}
		}
		return LJ_energy;
	}
	void Field_energy_system::energy_gradient_at(const double *const x, double *gradient)
	{
		molecular_energy_gradient_at(x, gradient);
	}
	void Field_energy_system::molecular_energy_gradient_at(const double *const x, double *gradient)
	{
		// gradient of distance energy 
		std::vector<Packing_polygon> transformed_polygons;
		transformed_polygons.reserve(polygons.size());
		for (unsigned int i = 0; i < polygons.size(); i++)
		{
			transformed_polygons.push_back(Packing_polygon(polygons[i]->vertices_begin(), polygons[i]->vertices_end()));
			transformed_polygons.back().transform(Parameter(1.0, x[i*3], x[i*3+1], x[i*3+2]));
			gradient[i*3] = gradient[i*3+1] = gradient[i*3+2] = 0.0;
		}

		for (unsigned int i = 0; i < transformed_polygons.size(); i++)
		{
			double si = std::sin(x[3*i]), ci = std::cos(x[3*i]);
			for (unsigned int j = 0; j < transformed_polygons.size(); j++)
			{
				if (j == i) continue;
				double sj = std::sin(x[3*j]), cj = std::cos(x[3*j]);
				for (unsigned int k = 0; k < transformed_polygons[i].size(); k++)
				{
					unsigned int edge_idx;
					double d2 = point_polygon_squared_distance(transformed_polygons[i].vertex(k), transformed_polygons[j], edge_idx);
					Vector_2 n = polygons[j]->edge(edge_idx).to_vector().perpendicular(CGAL::COUNTERCLOCKWISE);
					// take an arbitrary point(middle here) and the normal vector to describe this segment
					n = n/CGAL::sqrt(n.squared_length());
					Point_2 p = CGAL::midpoint(polygons[j]->edge(edge_idx).source(), polygons[j]->edge(edge_idx).target());
					Transformation_2 Ri(ci, -si, 0.0, si, ci, 0.0), Rj(cj, -sj, 0.0, sj, cj, 0.0);
					Vector_2 v1 = Rj.transform(n); // normal vector after transformation and line
					Vector_2 v2 = Ri.transform(polygons[i]->vertex(k)) + Vector_2(x[3*i+1], x[3*i+2]) - Rj.transform(p) - Vector_2(x[3*j+1], x[3*j+2]);
					double m = v1*v2;

					//double dist_grad_x = 2*m*v1.x(), dist_grad_y = 2*m*v1.y();
					double rm2_over_d2;
					//if (d2 == 0.0) 
					//	rm2_over_d2 = 0.0;
					//else
						rm2_over_d2 = rm*rm/d2;
					double rm6_over_d6 = std::pow(rm2_over_d2, 3.0);
					double common_factor = rm6_over_d6*(rm6_over_d6 - 1);
					common_factor /= d2;
					
					gradient[3*i+1] += common_factor*(2*m*v1.x());
					gradient[3*i+2] += common_factor*(2*m*v1.y());
					Transformation_2 dRi_dthetai(-si, -ci, 0.0, ci, -si, 0.0), dRj_dthetaj(-sj, -cj, 0.0, cj, -sj, 0.0);
					//double tmp = dist_grad_x*(-polygons[i]->vertex(v).x()*s-polygons[i]->vertex(v).y()*c)+dist_grad_y*(polygons[i]->vertex(v).x()*c-polygons[i]->vertex(v).y()*s);
					double tmp = v1*dRi_dthetai(Vector_2(CGAL::ORIGIN, polygons[i]->vertex(k)))*(2*m);
					gradient[3*i] += tmp*common_factor;
					//gradient[3*j] += -gradient[3*i];
					gradient[3*j+1] += -gradient[3*i+1];
					gradient[3*j+2] += -gradient[3*i+2];
					tmp = dRj_dthetaj(n)*Ri(Vector_2(CGAL::ORIGIN, polygons[i]->vertex(k))) + dRj_dthetaj(n)*Vector_2(x[3*i+1]-x[3*j+1], x[3*i+2]-x[3*j+2]);
					tmp *= 2*m;
					gradient[3*j] += tmp*common_factor;
				}
			}
			gradient[3*i] *= -6;
			gradient[3*i+1] *= -6;
			gradient[3*i+2] *= -6;
		}
	}
	void Field_energy_system::constraint_at(const double *const x, double *constraint_values)
	{
		std::vector<Packing_polygon> transformed_polygons;
		transformed_polygons.reserve(polygons.size());
		for (unsigned int i = 0; i < polygons.size(); i++)
		{
			transformed_polygons.push_back(Packing_polygon(polygons[i]->vertices_begin(), polygons[i]->vertices_end()));
			transformed_polygons.back().transform(Parameter(1.0, x[i*3], x[i*3+1], x[i*3+2]));
		}
		for (unsigned int i = 0; i < transformed_polygons.size(); i++)
		{
			for (unsigned int j = 0; j < transformed_polygons[i].size(); j++)
			{
				*constraint_values++ = transformed_polygons[i].vertex(j).x();
				*constraint_values++ = transformed_polygons[i].vertex(j).y();
			}
		}
	}
	void Field_energy_system::constraint_grad_at(const double *const x, double *constraint_grad)
	{
		std::vector<Packing_polygon> transformed_polygons;
		transformed_polygons.reserve(polygons.size());
		for (unsigned int i = 0; i < polygons.size(); i++)
		{
			transformed_polygons.push_back(Packing_polygon(polygons[i]->vertices_begin(), polygons[i]->vertices_end()));
			transformed_polygons.back().transform(Parameter(1.0, x[i*3], x[i*3+1], x[i*3+2]));
		}
		for (unsigned int i = 0; i < transformed_polygons.size(); i++)
		{
			double fsin = std::sin(x[3*i]), fcos = std::cos(x[3*i]);
			for (unsigned int j = 0; j < transformed_polygons[i].size(); j++)
			{
				Packing_polygon::Point_2 c = transformed_polygons[i].vertex(j);
				*constraint_grad++ = -c.x()*fsin - c.y()*fcos;
				*constraint_grad++ = 1.0;
				*constraint_grad++ = 0.0;
				*constraint_grad++ = c.x()*fcos - c.y()*fsin;
				*constraint_grad++ = 0.0;
				*constraint_grad++ = 1.0;
			}
		}
	}
	int Field_energy_system::optimize(double *trans, bool apply_constraints)
	{
		nonoverlap_constraint = apply_constraints;
		int  nStatus;
		/* variables that are passed to KNITRO */
		KTR_context *kc;
		int n, m, nnzJ, nnzH, objGoal, objType;
		int *cType, *jacIndexVars, *jacIndexCons;
		double obj, *lambda;
		double *xLoBnds, *xUpBnds, *xInitial, *cLoBnds, *cUpBnds;

		n = polygons.size()*3;
		m = 0;
		for (unsigned int i = 0; i < polygons.size(); i++)
			m += polygons[i]->size();
		m *= 2; // constraint on the x and y coordinate
		nnzJ = m*3;
		nnzH = 0;
		lambda = (double *)malloc((m+n) * sizeof(double));
		xLoBnds = (double *)malloc(n * sizeof(double));
		xUpBnds = (double *)malloc(n * sizeof(double));
		xInitial = (double *)malloc(n * sizeof(double));
		cType        = (int    *)malloc(m * sizeof(int));
		cLoBnds      = (double *)malloc(m * sizeof(double));
		cUpBnds      = (double *)malloc(m * sizeof(double));
		jacIndexVars = (int    *)malloc(nnzJ * sizeof(int));
		jacIndexCons = (int    *)malloc(nnzJ * sizeof(int));
		/* objective type */
		objType = KTR_OBJTYPE_GENERAL;
		objGoal = KTR_OBJGOAL_MINIMIZE;
		/* bounds and constraints type */
		for (unsigned int i = 0; i < polygons.size(); i++) 
		{
			xLoBnds[3*i] = -boost::math::constants::pi<double>()/3;
			xUpBnds[3*i] = boost::math::constants::pi<double>()/3;
			xLoBnds[3*i+1] = xLoBnds[3*i+2] = -KTR_INFBOUND;
			xUpBnds[3*i+1] = xUpBnds[3*i+2] = KTR_INFBOUND;
		}
		int *ptrCType = cType;
		double *ptrCLoBnds = cLoBnds, *ptrCUpBnds = cUpBnds;
		for (unsigned int i = 0; i < polygons.size(); i++)
		{
			for (unsigned int j = 0; j < polygons[i]->size(); j++)
			{
				*ptrCType++ = KTR_CONTYPE_GENERAL;
				*ptrCLoBnds++ = bounding_box.x_min();
				*ptrCUpBnds++ = bounding_box.x_max();
				*ptrCType++ = KTR_CONTYPE_GENERAL;
				*ptrCLoBnds++ = bounding_box.y_min();
				*ptrCUpBnds++ = bounding_box.y_max();				
			}
		}
		int constraint_no = 0;
		int *ptrCons = jacIndexCons, *ptrVars = jacIndexVars;
		for (unsigned int i = 0; i < polygons.size(); i++)
		{
			for (unsigned int j = 0; j < polygons[i]->size(); j++)
			{
				*ptrCons++ = constraint_no;
				*ptrVars++ = 3*i;
				*ptrCons++ = constraint_no;
				*ptrVars++ = 3*i+1;
				*ptrCons++ = constraint_no;
				*ptrVars++ = 3*i+2;
				constraint_no++;
				/* y constraint for this circle */
				*ptrCons++ = constraint_no;
				*ptrVars++ = 3*i;
				*ptrCons++ = constraint_no;
				*ptrVars++ = 3*i+1;
				*ptrCons++ = constraint_no;
				*ptrVars++ = 3*i+2;
				constraint_no++;
			}
		}
		std::copy(trans, trans+n, xInitial);
		/* create a KNITRO instance */
		kc = KTR_new();
		if (kc == NULL)
			return -1000 ; // probably a license issue
		if (KTR_set_int_param_by_name(kc, "algorithm", 0) != 0)
			return -1000;
		/* set options: automatic gradient and hessian matrix */
		//if (KTR_set_int_param_by_name(kc, "bar_murule", 2) != 0)
		//	return -1000;
		if (KTR_set_int_param_by_name(kc, "gradopt", KTR_GRADOPT_EXACT) != 0)
			return -1000 ;
		if (KTR_set_int_param_by_name(kc, "hessopt", KTR_HESSOPT_BFGS) != 0)
			return -1000 ;
		if (KTR_set_int_param_by_name(kc, "outlev", 1) != 0)
			return -1000 ;
		if (KTR_set_int_param_by_name(kc, "outmode", 0) != 0)
			return -1000;
		if (KTR_set_int_param_by_name(kc, "maxit", 0) != 0)
			return -1000 ;
		//if (KTR_set_int_param_by_name(kc, "ms_enable", 1) != 0)
		//	return -1000;
		/* register the callback function */
		if (KTR_set_func_callback(kc, &callback) != 0)
			return -1000 ;
		if (KTR_set_grad_callback(kc, &callback) != 0)
			return -1000 ;
		// the following two setttings are for debug
		//if (KTR_set_int_param_by_name(kc, "newpoint", 3) != 0)
		//	return -1000;
		//if (KTR_set_newpoint_callback(kc, &callback) !=0 )
		//	return -1000;
		/* pass the problem definition to KNITRO */
		nStatus = KTR_init_problem(kc, n, objGoal, objType, xLoBnds, xUpBnds, m, cType, cLoBnds, cUpBnds,
									nnzJ, jacIndexVars, jacIndexCons, nnzH, NULL, NULL, xInitial, NULL);
		/* free memory (KNITRO maintains its own copy) */
		free(xLoBnds);	free(xUpBnds); free(xInitial);
		free(cType); free(cLoBnds);	free(cUpBnds);
		free(jacIndexVars);	free(jacIndexCons);

		nStatus = KTR_solve(kc, trans, lambda, 0, &obj, NULL, NULL, NULL, NULL, NULL, this);

		/* delete the KNITRO instance and primal/dual solution */
		KTR_free(&kc);
		free(lambda);
		return nStatus;
	}
}